from alibabacloud_tea_openapi.models import Config
from alibabacloud_searchplat20240529.client import Client
from alibabacloud_searchplat20240529.models import GetTextEmbeddingRequest
import os
import numpy as np
import pandas as pd


def find_md_files(directory):
    md_files = []
    for root, dirs, files in os.walk(directory):
        for file in files:
            if file.endswith('.md'):
                # 获取文件的绝对路径
                abs_path = os.path.abspath(os.path.join(root, file))
                md_files.append(abs_path)
    return md_files

def cosine_similarity(A, B):
    dot_product = np.dot(A, B)
    norm_A = np.linalg.norm(A)
    norm_B = np.linalg.norm(B)
    similarity = dot_product / (norm_A * norm_B)
    return similarity


if __name__ == '__main__':
    md_files = find_md_files('static/doc/split')
    embeddings = {}
    # token配置，endpoint配置
    config = Config(bearer_token="OS-g1h6d9g3s948p1nu1",
                    endpoint="default-dm5.platform-cn-shanghai.opensearch.aliyuncs.com",
                    protocol="http")
    client = Client(config=config)
    for md_f in md_files:
        with open(md_f, 'r') as file:
            input = file.read()
            request = GetTextEmbeddingRequest(input=[input], input_type="document")
            response = client.get_text_embedding("default", "ops-text-embedding-001", request)
            res = response.body.result.embeddings
            print(res[0])
            embeddings[os.path.basename(md_f)] = np.array(res[0].embedding)

    query = '如何加强电力产品增值税的征收管理'
    request = GetTextEmbeddingRequest(input=[query], input_type="query")
    response = client.get_text_embedding("default", "ops-text-embedding-001", request)
    query_embeddings = np.array(response.body.result.embeddings[0].embedding)

    column_names = ['file', 'similarity']
    df = pd.DataFrame(columns=column_names)
    for file_path, doc_embedding in embeddings.items():
        similarity = cosine_similarity(query_embeddings, doc_embedding)
        df.loc[len(df)] = {"file": file_path, "similarity": similarity}

    df = df.sort_values(by='similarity', ascending=False)
    top1 = df.head(1)
    top5 = df.head(5)
    top10 = df.head(10)
    print("top1")
    print(top1)
    print("top5")
    print(top5)
    print("top10")
    print(top10)

    # df['label'] = 1 if df['file'] == '10.md' else 0
    df['label'] = df['file'].apply(lambda x: 1 if x == '10.md' or x == '09.md' else 0)

    from sklearn.metrics import average_precision_score
    def calculate_map(df, ground_truth_column, prediction_column):
        y_true = df[ground_truth_column]
        y_scores = df[prediction_column]

        map_value = average_precision_score(y_true, y_scores)
        return map_value


    map_value = calculate_map(df, 'label', 'similarity')
    print('mAP:', map_value)


